package cn.lg.soar.mvc.filter;

import jakarta.servlet.*;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.Assert;

import java.io.IOException;
import java.util.*;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

import static cn.lg.soar.mvc.model.RefererFilterRule.*;

/**
 * Referer 黑名单
 * @author luguoxiang
 * @date 2022/9/5
 * 开源项目：https://gitee.com/lgx1992/lg-soar 求star！请给我star！请帮我点个star！
 */
public class RefererBlacklistFilter implements Filter {

    private static final Logger logger = LoggerFactory.getLogger(RefererBlacklistFilter.class);

    private final Set<String> eqs;
    private final String[] starts;
    private final String[] ends;
    private final Pattern[] regexs;
    private final boolean nullable;

    public RefererBlacklistFilter(Collection<String> referers, boolean nullable) {
        Assert.notEmpty(referers, "referers must not be empty");
        this.nullable = nullable;
        int size = referers.size();
        List<String> eqs = new ArrayList<>(size);
        List<String> starts = new ArrayList<>(size);
        List<String> ends = new ArrayList<>(size);
        List<String> regexs = new ArrayList<>(size);
        referers.stream().map(String::toLowerCase).distinct().forEach(x -> {
            if (x.startsWith(REGEX)) {
                regexs.add(x.substring(REGEX.length()));
                return;
            }
            if (x.startsWith(EQ)) {
                eqs.add(x.substring(EQ.length()));
                return;
            }
            if (x.startsWith(END)) {
                ends.add(x.substring(END.length()));
                return;
            }
            starts.add(x);
        });
        if (eqs.isEmpty()) {
            this.eqs = null;
        } else {
            this.eqs = new HashSet<>(eqs);
        }
        if (starts.isEmpty()) {
            this.starts = null;
        } else {
            this.starts = starts.toArray(new String[0]);
        }
        if (ends.isEmpty()) {
            this.ends = null;
        } else {
            this.ends = ends.toArray(new String[0]);
        }
        if (regexs.isEmpty()) {
            this.regexs = null;
        } else {
            this.regexs = regexs.stream().map(Pattern::compile).toArray(Pattern[]::new);
        }
    }

    @Override
    public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain)
            throws IOException, ServletException {

        HttpServletRequest httpRequest = (HttpServletRequest) request;
        HttpServletResponse httpResponse = (HttpServletResponse) response;

        String referer = httpRequest.getHeader("Referer");

        // 允许空 Referer？
        if (referer == null || referer.isEmpty()) {
            if (nullable) {
                chain.doFilter(request, response);
            } else {
                httpResponse.sendError(HttpServletResponse.SC_FORBIDDEN, "Empty Referer is not allowed.");
                logger.warn("Empty Referer is not allowed.");
            }
            return;
        }
        referer = referer.toLowerCase();
        if (check(referer)) {
            chain.doFilter(request, response);
        } else {
            httpResponse.sendError(HttpServletResponse.SC_FORBIDDEN, "Referer not allowed: " + referer);
            logger.warn("Referer not allowed: {}", referer);
        }

    }

    public boolean check(String referer) {
        if (eqs != null) {
            if (eqs.contains(referer)) {
                return false;
            }
        }
        if (starts != null) {
            for (String x : starts) {
                if (referer.startsWith(x)) {
                    return false;
                }
            }
        }
        if (ends != null) {
            for (String x : ends) {
                if (referer.endsWith(x)) {
                    return false;
                }
            }
        }
        if (regexs != null) {
            for (Pattern regex : regexs) {
                if (regex.matcher(referer).matches()) {
                    return false;
                }
            }
        }
        return true;
    }

}